package com.de.rtsp;

import com.de.controller.IndexController;
import com.de.entity.Image;
import com.de.service.WebSocketServer;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.bytedeco.ffmpeg.global.avutil;
import org.bytedeco.javacv.*;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.imageio.IIOImage;
import javax.imageio.ImageIO;
import javax.imageio.ImageWriteParam;
import javax.imageio.ImageWriter;
import javax.imageio.stream.ImageOutputStream;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.Iterator;

import static org.bytedeco.ffmpeg.global.avutil.AV_LOG_ERROR;

/**
 * 转换rtsp为img
 *
 * @author IT_CREATE
 * @date 2021/6/8 12:00:00
 */
@Slf4j
@Component
@Data
public class MediaTransfer {

    private String rtspUrl;

    @Value("${rtsp.transport.type}")
    private String rtspTransportType;

    private String readOrWriteTimeout = "5000000";
    private String netTimeout = "5000000";

    private FFmpegFrameGrabber grabber;
    private volatile boolean isStreaming = false;
    private Thread streamThread;

    public String getRtspUrl() {
        return rtspUrl;
    }

    public void setRtspUrl(String rtspUrl) {
        this.rtspUrl = rtspUrl;
    }

    public int frameRate = 120;
    public int frameWidth = 960;
    public int frameHeight = 540;

    // Timeout settings
    private long timeoutMillis = 15000; // 15 seconds timeout

    private static final int MAX_NULL_FRAME_THRESHOLD = 10000; // 5 seconds
    private static final int MAX_NULL_SOCKET_THRESHOLD = 10000; // 5 seconds
    private static final int SLEEP_DELAY = 10000; // 1 second
    private static final int MIN_LOG_COUNT = 5;

    public void startStreaming(String id) {
        if (!isStreaming) {
            log.info("Starting RTSP stream...");
            grabber = createGrabber(rtspUrl);
            if (grabber != null) {
                isStreaming = true;
                streamThread = new Thread(() -> startCameraPush(id));
                streamThread.start();
            } else {
                log.error("Failed to create grabber.");
            }
        } else {
            log.warn("Streaming is already in progress.");
        }
    }

    public FFmpegFrameGrabber createGrabber(String rtsp) {
        try {
            grabber = FFmpegFrameGrabber.createDefault(rtsp);
            grabber.setOption("stimeout", netTimeout);
            grabber.setOption("skip_frame", "nokey");
            grabber.setOption("threads", "1");
            grabber.setOption("buffer_size", "1024000");
            grabber.setOption("rw_timeout", readOrWriteTimeout);
            grabber.setOption("rtsp_transport", rtspTransportType);
            grabber.setFrameRate(frameRate);
            grabber.setImageWidth(frameWidth);
            grabber.setImageHeight(frameHeight);
            grabber.setVideoBitrate(1000000);
            avutil.av_log_set_level(AV_LOG_ERROR);
            grabber.setOption("color_range", "limited");

            // 添加以下配置
            grabber.setOption("max_delay", "500000"); // 设置最大延迟
            grabber.setOption("fflags", "nobuffer"); // 禁用帧缓冲
            grabber.setOption("flags", "low_delay"); // 低延迟模式
            grabber.setOption("probesize", "32"); // 减小探测大小
            grabber.setOption("analyzeduration", "0"); // 减小分析时间

            grabber.setOption("r", String.valueOf(frameRate)); // 强制输入帧率

            // 强制使用TCP传输，可以提高稳定性
            grabber.setOption("rtsp_transport", rtspTransportType);

            // 设置更激进的超时参数
            grabber.setOption("stimeout", "5000000");
            grabber.setOption("rw_timeout", "5000000");

            return grabber;
        } catch (FrameGrabber.Exception e) {
            log.error("Failed to create FFmpegFrameGrabber", e);
            return null;
        }
    }

    public void startCameraPush(String id) {
        Java2DFrameConverter java2DFrameConverter = new Java2DFrameConverter();
        StreamingMetrics metrics = new StreamingMetrics();
        Frame lastValidFrame = null; // 保存上一个有效帧

        try {
            grabber.start();
            log.info("Grabber 启动成功。");

            // 优化帧率控制
            long frameInterval = 1000 / frameRate; // 计算帧间隔
            long lastFrameTime = System.currentTimeMillis();
            int consecutiveAbnormalFrames = 0; // 连续异常帧计数

            while (isStreaming) {
                if (!handleSocketConnection(id, metrics)) {
                    continue;
                }

                // 严格的帧率控制
                long currentTime = System.currentTimeMillis();
                long elapsedTime = currentTime - lastFrameTime;
                if (elapsedTime < frameInterval) {
                    Thread.sleep(frameInterval - elapsedTime);
                    continue;
                }

                Frame frame = grabber.grabImage();
                if (!handleFrame(frame, id, metrics)) {
                    // 如果当前帧异常且有上一个有效帧，使用上一个有效帧
                    if (lastValidFrame != null && consecutiveAbnormalFrames < 5) {
                        frame = lastValidFrame;
                        consecutiveAbnormalFrames++;
                    } else {
                        continue;
                    }
                } else {
                    consecutiveAbnormalFrames = 0;
                    lastValidFrame = frame; // 更新最后的有效帧
                }

                processFrame(frame, java2DFrameConverter, id);
                lastFrameTime = currentTime;
            }
        } catch (FrameGrabber.Exception e) {
            log.error("抓取器出错，尝试重新连接", e);
            // 发生错误时尝试重新连接
            reconnect(id);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        } finally {
            cleanup(id);
        }
    }

    private boolean handleSocketConnection(String id, StreamingMetrics metrics)
            throws InterruptedException {
        if (WebSocketServer.getOnlineCount() == 0) {
            metrics.incrementNullSocketCount();

            if (metrics.shouldLogSocketWarning()) {
                log.warn("接收到空连接，计数: " + metrics.getNullSocketCount());
            }

            if (metrics.hasExceededSocketTimeout()) {
                log.warn("在 {}ms 内没有检测到Socket连接。停止流媒体传输。id:{}",
                    MAX_NULL_SOCKET_THRESHOLD, id);
                stopStreaming();
                IndexController.mediaTransfers.remove(id);
                return false;
            }
            Thread.sleep(SLEEP_DELAY);
            return false;
        }
        metrics.resetSocketMetrics();
        return true;
    }

    private boolean handleFrame(Frame frame, String id, StreamingMetrics metrics)
            throws InterruptedException {
        if (frame == null) {
            metrics.incrementNullFrameCount();

            if (metrics.shouldLogFrameWarning()) {
                log.warn("接收到空帧，计数: " + metrics.getNullFrameCount());
            }

            if (metrics.hasExceededFrameTimeout()) {
                log.warn("在 {}ms 内未接收到有效帧。停止流媒体传输。id:{}",
                    MAX_NULL_FRAME_THRESHOLD, id);
                stopStreaming();
                IndexController.mediaTransfers.remove(id);
                return false;
            }
            Thread.sleep(SLEEP_DELAY);
            return false;
        }
        metrics.resetFrameMetrics();
        return true;
    }

    private void processFrame(Frame frame, Java2DFrameConverter converter, String id) throws InterruptedException {
        if (frame == null || frame.image == null) {
            log.debug("跳过空帧或无图像数据的帧");
            return;
        }

        BufferedImage bufferedImage = converter.getBufferedImage(frame);
        if (bufferedImage == null) {
            log.debug("无效图像");
            return;
        }

        // 增加多重图像质量检查
        if (isAbnormalFrame(bufferedImage)) {
            log.debug("检测到异常帧，已跳过");
            return;
        }

        byte[] bytes = compressImage(bufferedImage, "jpg", 0.8f);
        if (bytes != null && bytes.length > 0) {
            WebSocketServer.sendInfo(new Image(id, bytes), id);
        }
    }

    private boolean isAbnormalFrame(BufferedImage image) {
        return isGrayFrame(image) ||
               isDarkFrame(image) ||
               isOverExposedFrame(image) ||
               isNoisyFrame(image);
    }

    private boolean isGrayFrame(BufferedImage image) {
        if (image == null) return true;

        int width = image.getWidth();
        int height = image.getHeight();
        int sampleSize = 20; // 增加采样点数量
        int tolerance = 15;  // 增加容差值

        for (int i = 0; i < sampleSize; i++) {
            int x = (width * (i + 1)) / (sampleSize + 1);
            for (int j = 0; j < sampleSize; j++) {
                int y = (height * (j + 1)) / (sampleSize + 1);
                int rgb = image.getRGB(x, y);
                int r = (rgb >> 16) & 0xFF;
                int g = (rgb >> 8) & 0xFF;
                int b = rgb & 0xFF;

                if (Math.abs(r - g) > tolerance ||
                    Math.abs(g - b) > tolerance ||
                    Math.abs(r - b) > tolerance) {
                    return false;
                }
            }
        }
        return true;
    }

    private boolean isDarkFrame(BufferedImage image) {
        int width = image.getWidth();
        int height = image.getHeight();
        int sampleSize = 20;
        int darkThreshold = 30; // 暗度阈值
        int darkPixelCount = 0;
        int totalSamples = sampleSize * sampleSize;

        for (int i = 0; i < sampleSize; i++) {
            int x = (width * (i + 1)) / (sampleSize + 1);
            for (int j = 0; j < sampleSize; j++) {
                int y = (height * (j + 1)) / (sampleSize + 1);
                int rgb = image.getRGB(x, y);
                int brightness = (((rgb >> 16) & 0xFF) +
                                ((rgb >> 8) & 0xFF) +
                                (rgb & 0xFF)) / 3;
                if (brightness < darkThreshold) {
                    darkPixelCount++;
                }
            }
        }

        return (darkPixelCount / (double)totalSamples) > 0.8; // 80%以上像素过暗判定为异常
    }

    private boolean isOverExposedFrame(BufferedImage image) {
        int width = image.getWidth();
        int height = image.getHeight();
        int sampleSize = 20;
        int brightThreshold = 240; // 过曝阈值
        int brightPixelCount = 0;
        int totalSamples = sampleSize * sampleSize;

        for (int i = 0; i < sampleSize; i++) {
            int x = (width * (i + 1)) / (sampleSize + 1);
            for (int j = 0; j < sampleSize; j++) {
                int y = (height * (j + 1)) / (sampleSize + 1);
                int rgb = image.getRGB(x, y);
                int brightness = (((rgb >> 16) & 0xFF) +
                                ((rgb >> 8) & 0xFF) +
                                (rgb & 0xFF)) / 3;
                if (brightness > brightThreshold) {
                    brightPixelCount++;
                }
            }
        }

        return (brightPixelCount / (double)totalSamples) > 0.5; // 50%以上像素过亮判定为异常
    }

    private boolean isNoisyFrame(BufferedImage image) {
        int width = image.getWidth();
        int height = image.getHeight();
        int sampleSize = 20;
        int noiseThreshold = 80; // 增加噪点阈值，更宽容一些
        int noisyPixelCount = 0;
        int totalSamples = (sampleSize - 1) * (sampleSize - 1);

        // 使用滑动窗口检测噪点
        for (int i = 1; i < width - 1; i += width/sampleSize) {
            for (int j = 1; j < height - 1; j += height/sampleSize) {
                int centerPixel = image.getRGB(i, j);
                int[] neighbors = {
                    image.getRGB(i-1, j),
                    image.getRGB(i+1, j),
                    image.getRGB(i, j-1),
                    image.getRGB(i, j+1)
                };

                // 计算与周围像素的平均差异
                int avgDiff = 0;
                for (int neighbor : neighbors) {
                    avgDiff += getPixelDifference(centerPixel, neighbor);
                }
                avgDiff /= 4;

                if (avgDiff > noiseThreshold) {
                    noisyPixelCount++;
                }
            }
        }

        return (noisyPixelCount / (double)totalSamples) > 0.4; // 提高容忍度
    }

    private int getPixelDifference(int rgb1, int rgb2) {
        int r1 = (rgb1 >> 16) & 0xFF;
        int g1 = (rgb1 >> 8) & 0xFF;
        int b1 = rgb1 & 0xFF;

        int r2 = (rgb2 >> 16) & 0xFF;
        int g2 = (rgb2 >> 8) & 0xFF;
        int b2 = rgb2 & 0xFF;

        return Math.abs(r1 - r2) + Math.abs(g1 - g2) + Math.abs(b1 - b2);
    }

    private void cleanup(String id) {
        IndexController.mediaTransfers.remove(id);
        stopGrabber();
    }

    // 新增内部类用于跟踪指标
    private static class StreamingMetrics {
        private int nullFrameCount = 0;
        private long startNullFrameTime = 0;
        private int nullSocketCount = 0;
        private long startNullSocketTime = 0;

        public void incrementNullFrameCount() {
            nullFrameCount++;
            if (nullFrameCount == 1) {
                startNullFrameTime = System.currentTimeMillis();
            }
        }

        public void incrementNullSocketCount() {
            nullSocketCount++;
            if (nullSocketCount == 1) {
                startNullSocketTime = System.currentTimeMillis();
            }
        }

        public boolean shouldLogFrameWarning() {
            return nullFrameCount % MIN_LOG_COUNT == 0;
        }

        public boolean shouldLogSocketWarning() {
            return nullSocketCount % MIN_LOG_COUNT == 0;
        }

        public boolean hasExceededFrameTimeout() {
            return System.currentTimeMillis() - startNullFrameTime > MAX_NULL_FRAME_THRESHOLD;
        }

        public boolean hasExceededSocketTimeout() {
            return System.currentTimeMillis() - startNullSocketTime > MAX_NULL_SOCKET_THRESHOLD;
        }

        public void resetFrameMetrics() {
            nullFrameCount = 0;
            startNullFrameTime = 0;
        }

        public void resetSocketMetrics() {
            nullSocketCount = 0;
            startNullSocketTime = 0;
        }

        public int getNullFrameCount() {
            return nullFrameCount;
        }

        public int getNullSocketCount() {
            return nullSocketCount;
        }
    }

    public void stopStreaming() {
        isStreaming = false; // Stop the streaming loop
        if (streamThread != null) {
            streamThread.interrupt(); // Interrupt the thread to stop blocking calls
            try {
                streamThread.join(); // Wait for the thread to finish
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt(); // Restore interrupted status
            }
        }
        stopGrabber(); // Ensure the grabber is stopped
        log.info("Streaming stopped.");
    }

    public void stopGrabber() {
        if (grabber != null) {
            try {
                grabber.close();
                log.info("Grabber stopped successfully.");
            } catch (FrameGrabber.Exception e) {
                log.error("Error stopping grabber", e);
            } catch (Exception e) {
                log.error("Unexpected error stopping grabber", e);
            } finally {
                this.grabber = null; // Clear grabber reference
            }
        }
    }

    public void stopGrabber(FFmpegFrameGrabber grabber) {
        if (grabber != null) {
            try {
                grabber.close();
                log.info("Grabber stopped successfully.");
            } catch (FrameGrabber.Exception e) {
                log.error("Error stopping grabber", e);
            } catch (Exception e) {
                log.error("Unexpected error stopping grabber", e);
            }
        }
    }

    private byte[] compressImage(BufferedImage image, String format, float quality) {
        try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
            Iterator<ImageWriter> writers = ImageIO.getImageWritersByFormatName(format);
            if (!writers.hasNext()) return null;

            ImageWriter writer = writers.next();
            ImageWriteParam params = writer.getDefaultWriteParam();
            if (params.canWriteCompressed()) {
                params.setCompressionMode(ImageWriteParam.MODE_EXPLICIT);
                params.setCompressionQuality(quality);
            }

            ImageOutputStream ios = ImageIO.createImageOutputStream(out);
            writer.setOutput(ios);
            writer.write(null, new IIOImage(image, null, null), params);
            writer.dispose();
            return out.toByteArray();
        } catch (IOException e) {
            log.error("Error compressing image", e);
            return null;
        }
    }

    // 添加重连方法
    private void reconnect(String id) {
        stopGrabber();
        try {
            Thread.sleep(2000); // 等待2秒后重试
            startStreaming(id);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }

}
